import torch.nn as nn


class ResDec(nn.Module):
    def __init__(self, in_dementions, out_dementions, kernel_size=1, padding=0, stride=1, out_padding=0):
        super().__init__()
        self.conv0 = nn.Sequential(nn.Conv2d(in_dementions, in_dementions, kernel_size=1),
                                   nn.BatchNorm2d(in_dementions),
                                   nn.ReLU())
        self.conv_t1 = nn.Sequential(
            nn.ConvTranspose2d(in_dementions, in_dementions, kernel_size=kernel_size, padding=padding, stride=stride,
                               output_padding=out_padding),
            nn.BatchNorm2d(in_dementions),
            nn.ReLU())
        self.conv2 = nn.Sequential(nn.Conv2d(in_dementions, out_dementions, kernel_size=1),
                                   nn.BatchNorm2d(out_dementions),
                                   nn.ReLU())

    def forward(self, x):
        out = self.conv2(self.conv_t1(self.conv0(x)))
        return out


class Decoder(nn.Module):
    def __init__(self, filters, output_dementions=1, bias=True, name='decoder'):
        super().__init__()
        self.bias = bias
        self.name = name

        self.dec4 = ResDec(filters[4], filters[3], stride=2, out_padding=1)
        self.dec3 = ResDec(filters[3], filters[2], kernel_size=3, padding=1, stride=2, out_padding=1)
        self.dec2 = ResDec(filters[2], filters[1], kernel_size=3, padding=1, stride=2, out_padding=1)
        self.dec1 = ResDec(filters[1], filters[0], kernel_size=3, padding=1)
        self.conv0 = nn.Conv2d(filters[0], output_dementions, kernel_size=1)
        # self.dec4 = ResDec(filters[4],filters[3],stride=2,out_padding=1)
        # self.dec3 = ResDec(filters[3], filters[2],kernel_size=1,stride=2,out_padding=1)
        # self.dec2 = ResDec(filters[2], filters[1],kernel_size=3,padding=1,stride=2,out_padding=1)
        # self.dec1 = ResDec(filters[1], filters[0],kernel_size=1)
        # self.conv0 = nn.Conv2d(filters[0],output_dementions,kernel_size=1)

    def forward(self, enc_layers):
        x, enc1, enc2, enc3, enc4, layer5 = enc_layers

        layer4 = self.dec4(layer5)
        layer4 += enc4
        layer3 = self.dec3(layer4)
        layer3 += enc3
        layer2 = self.dec2(layer3)
        layer2 += enc2
        layer1 = self.dec1(layer2)
        layer1 += enc1
        out = self.conv0(layer1)
        out += x[:, :1, :, :]

        return out
